import os
import torchvision
from typing import Any, Callable, Dict, List, Optional, Tuple


ROOT = f"e:/sfxData/DeepLearning"

# KINDS
FASHION_MNIST_KINDS = [
    't-shirt',
    'trouser',
    'pullover',
    'dress',
    'coat',
    'sandal',
    'shirt',
    'sneaker',
    'bag',
    'ankle boot']

FASHION_MNIST_TRANSFORM_DEFAULT = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),                      # 转换为张量
    # torchvision.transforms.Normalize((0.1307,), (0.3081,))  # 设定标准化值
])


def get_fashion_mnist(root: str, train: bool = True, transform: Optional[Callable] = None) -> torchvision.datasets.FashionMNIST:
    """
    取得 FashionMNIST 数据集.

    参数:
        root (string): 数据集根目录
        train (bool, optional): 是否取得训练集
        transform (callable, optional): 转换
    """
    xtransform = FASHION_MNIST_TRANSFORM_DEFAULT if (
        transform is None) else transform
    isDownload = not(os.path.exists(root)) or not os.listdir(root)
    dset = torchvision.datasets.FashionMNIST(
        root, train, xtransform, None, download=isDownload)
    return dset


def get_fashion_mnist_kind_name(index: int) -> str:
    return FASHION_MNIST_KINDS[index]


def get_fashion_mnist_labels(labels) -> list[str]:
    test_labels = ['t-shirt', 'trouser', 'pullover', 'dress',
                   'coat', 'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return [test_labels[int(i)] for i in labels]
